import numpy as np
from . import LocalSearchOperator, register_operator


class ThreeOptOperator(LocalSearchOperator):
    """通用的3-opt局部搜索算子"""

    def __init__(self):
        self.name = "three_opt"
        self.supported_problems = ["tsp"]

    def generate_neighborhood(self, solution, problem_instance):
        """生成3-opt邻域"""
        neighborhood = []
        n = len(solution)

        # 遍历所有可能的三个断点组合
        for i in range(1, n - 4):
            for j in range(i + 2, n - 2):
                for k in range(j + 2, n):
                    # 跳过无效的组合
                    if k - i <= 2 or j - i <= 1 or k - j <= 1:
                        continue

                    # 生成所有可能的3-opt移动
                    # 情况1: 反转中间段
                    new_solution1 = solution.copy()
                    new_solution1[i:j] = new_solution1[i:j][::-1]
                    neighborhood.append(new_solution1)

                    # 情况2: 反转后段
                    new_solution2 = solution.copy()
                    new_solution2[j:k] = new_solution2[j:k][::-1]
                    neighborhood.append(new_solution2)

                    # 情况3: 同时反转中间段和后段
                    new_solution3 = solution.copy()
                    new_solution3[i:j] = new_solution3[i:j][::-1]
                    new_solution3[j:k] = new_solution3[j:k][::-1]
                    neighborhood.append(new_solution3)

                    # 情况4: 交换中间段和后段
                    new_solution4 = solution.copy()
                    new_solution4[i:k] = np.concatenate([new_solution4[j:k], new_solution4[i:j]])
                    neighborhood.append(new_solution4)

                    # 情况5: 交换并反转中间段
                    new_solution5 = solution.copy()
                    new_solution5[i:k] = np.concatenate([new_solution4[j:k], new_solution4[i:j][::-1]])
                    neighborhood.append(new_solution5)

                    # 情况6: 交换并反转后段
                    new_solution6 = solution.copy()
                    new_solution6[i:k] = np.concatenate([new_solution4[j:k][::-1], new_solution4[i:j]])
                    neighborhood.append(new_solution6)

                    # 情况7: 交换并同时反转中间段和后段
                    new_solution7 = solution.copy()
                    new_solution7[i:k] = np.concatenate([new_solution4[j:k][::-1], new_solution4[i:j][::-1]])
                    neighborhood.append(new_solution7)

        return neighborhood

    def __call__(self, solution, problem_instance, intensity):
        """应用3-opt局部搜索"""
        if not self.is_applicable(problem_instance):
            return solution, 0

        best_solution = solution.copy()
        best_cost = problem_instance.evaluate(best_solution)
        improved = True
        iterations = 0

        while improved and iterations < intensity:
            improved = False
            neighborhood = self.generate_neighborhood(best_solution, problem_instance)

            for neighbor in neighborhood:
                neighbor_cost = problem_instance.evaluate(neighbor)

                if neighbor_cost < best_cost:
                    best_solution = neighbor
                    best_cost = neighbor_cost
                    improved = True
                    break

            iterations += 1

        improvement = problem_instance.evaluate(solution) - best_cost
        return best_solution, improvement


# 注册算子
register_operator("three_opt", ThreeOptOperator())